from PIL import Image
from PyQt6 import QtCore, QtWidgets
from PyQt6.QtCore import Qt

from qfluentwidgets import FolderListDialog, Dialog
from qfluentwidgets import FluentIcon as FIF
from PyQt6.QtGui import QPixmap
from model_predict import Handwritten_digit_recognition_Single, Handwritten_digit_recognition_coiled
from fenge import charSeperate
from qfluentwidgets import TogglePushButton, PixmapLabel

class Widget_hwdr(QtWidgets.QFrame):

    def __init__(self, text: str, parent=None):
        super().__init__(parent=parent)

        sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Policy.Preferred,
                                           QtWidgets.QSizePolicy.Policy.Expanding)
        sizePolicy.setHorizontalStretch(0)
        sizePolicy.setVerticalStretch(1)
        sizePolicy.setHeightForWidth(self.sizePolicy().hasHeightForWidth())
        self.setSizePolicy(sizePolicy)
        self.gridLayout = QtWidgets.QGridLayout(self)
        self.gridLayout.setObjectName("gridLayout")
        self.frame = QtWidgets.QFrame(parent=self)
        self.frame.setFrameShape(QtWidgets.QFrame.Shape.StyledPanel)
        self.frame.setFrameShadow(QtWidgets.QFrame.Shadow.Raised)
        self.frame.setObjectName("frame")
        self.horizontalLayout_2 = QtWidgets.QHBoxLayout(self.frame)
        self.horizontalLayout_2.setObjectName("horizontalLayout_2")

        self.pixmapLabel1 = PixmapLabel(self)
        self.pixmapLabel1.setObjectName("label")

        self.horizontalLayout_2.addWidget(self.pixmapLabel1)
        self.gridLayout.addWidget(self.frame, 0, 0, 1, 1)
        self.frame_2 = QtWidgets.QFrame(parent=self)
        self.frame_2.setFrameShape(QtWidgets.QFrame.Shape.StyledPanel)
        self.frame_2.setFrameShadow(QtWidgets.QFrame.Shadow.Raised)
        self.frame_2.setObjectName("frame_2")
        self.horizontalLayout = QtWidgets.QHBoxLayout(self.frame_2)
        self.horizontalLayout.setObjectName("horizontalLayout")
        self.toggleButton1 = TogglePushButton(FIF.FOLDER_ADD, 'Toggle push button', self)
        self.toggleButton1.clicked.connect(self.showDialog)
        self.horizontalLayout.addWidget(self.toggleButton1)
        self.toggleButton2 = TogglePushButton(FIF.PRINT, 'Toggle push button', self)
        self.toggleButton2.clicked.connect(self.showDialog_coiled)
        self.horizontalLayout.addWidget(self.toggleButton2)
        self.gridLayout.addWidget(self.frame_2, 0, 1, 1, 1)
        self.frame_3 = QtWidgets.QFrame(parent=self)
        self.frame_3.setFrameShape(QtWidgets.QFrame.Shape.StyledPanel)
        self.frame_3.setFrameShadow(QtWidgets.QFrame.Shadow.Raised)
        self.frame_3.setObjectName("frame_3")
        self.horizontalLayout_3 = QtWidgets.QHBoxLayout(self.frame_3)
        self.horizontalLayout_3.setObjectName("horizontalLayout_3")
        self.label_2 = QtWidgets.QLabel(parent=self.frame_3)
        self.label_2.setObjectName("label_2")
        self.horizontalLayout_3.addWidget(self.label_2)
        self.gridLayout.addWidget(self.frame_3, 1, 0, 1, 1)
        self.frame_4 = QtWidgets.QFrame(parent=self)
        self.frame_4.setFrameShape(QtWidgets.QFrame.Shape.StyledPanel)
        self.frame_4.setFrameShadow(QtWidgets.QFrame.Shadow.Raised)
        self.frame_4.setObjectName("frame_4")
        self.horizontalLayout_4 = QtWidgets.QHBoxLayout(self.frame_4)
        self.horizontalLayout_4.setObjectName("horizontalLayout_4")
        self.label_3 = QtWidgets.QLabel(parent=self.frame_4)
        self.label_3.setObjectName("label_3")
        self.horizontalLayout_4.addWidget(self.label_3)
        self.gridLayout.addWidget(self.frame_4, 1, 1, 1, 1)

        self.retranslateUi()
        QtCore.QMetaObject.connectSlotsByName(self)

    def retranslateUi(self):
        _translate = QtCore.QCoreApplication.translate
        self.setWindowTitle(_translate("Widget1", "Widget1"))
        self.pixmapLabel1.setText(_translate("Widget1", "原始图片展示区"))
        self.toggleButton1.setText(_translate("Widget1", "单个数字识别"))
        self.toggleButton2.setText(_translate("Widget1", "多个数字识别"))
        self.label_2.setText(_translate("Widget1", "预测结果"))
        self.label_3.setText(_translate("Widget1", "预测结果"))

    def showDialog(self):
        picture_paths = []
        title = '单个数字识别'
        content = "现在选择你要识别的图片"
        w = FolderListDialog(picture_paths, title, content, self)
        w.folderChanged.connect(self.Picture_settings_Single)  # 连接信号，并传递正确的方法
        w.exec()

    def showDialog_coiled(self):
        picture_paths = []
        title = '连续数字识别'
        content = "现在选择你要识别的图片"
        w = FolderListDialog(picture_paths, title, content, self)
        w.folderChanged.connect(self.Picture_settings_coiled)  # 连接信号，并传递正确的方法
        w.exec()

    def Picture_settings_Single(self, picture_paths):
        print(picture_paths)
        Predict_outcomes = Handwritten_digit_recognition_Single(picture_paths[0])
        for Predict_outcome in Predict_outcomes:
            im_array = Predict_outcome.plot()  # plot a BGR numpy array of predictions
            im = Image.fromarray(im_array[..., ::-1])  # RGB PIL image
            im.save('results.jpg')  # save image
            self.label_3.setText("预测结果是" + str(Predict_outcome.probs.top1))
        # 获取要显示图片的尺寸
        pixmap = QPixmap(picture_paths[0])
        self.pixmapLabel1.setPixmap(pixmap.scaled(self.pixmapLabel1.size()))
        self.pixmapLabel1.setScaledContents(True)
        self.pixmapLabel1.setMaximumSize(self.pixmapLabel1.width(), self.pixmapLabel1.height())

        result_pixmap = QPixmap("results.jpg")
        self.label_2.setPixmap(result_pixmap.scaled(self.label_2.size()))
        self.label_2.setScaledContents(True)

    def Picture_settings_coiled(self, picture_paths):
        img_list = charSeperate(picture_paths[0], '3x3')
        result = ""
        for img in img_list:
            Predict_outcomes = Handwritten_digit_recognition_coiled(img)
            for Predict_outcome in Predict_outcomes:
                result = result + str(Predict_outcome.probs.top1)
        pixmap = QPixmap(picture_paths[0])
        self.pixmapLabel1.setPixmap(pixmap.scaled(self.pixmapLabel1.size()))
        self.pixmapLabel1.setScaledContents(True)
        self.pixmapLabel1.setMaximumSize(self.pixmapLabel1.width(), self.pixmapLabel1.height())
        self.label_3.setText("预测结果是" + str("识别结果是  " + result))



